{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "9dc079e4",
   "metadata": {},
   "source": [
    "### Check SNP/indel variant count "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "8e626ec6",
   "metadata": {},
   "outputs": [],
   "source": [
    "from pathlib import Path\n",
    "import pandas as pd\n",
    "from pybedtools import BedTool\n",
    "from io import StringIO\n",
    "import pysam"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "f0a6be9d",
   "metadata": {},
   "outputs": [],
   "source": [
    "wkdir = \"/lustre/scratch126/humgen/projects/interval_rna/interval_rna_seq/thomasVDS/misexpression_v3/\"\n",
    "wkdir_path = Path(wkdir)\n",
    "\n",
    "chrom = \"chr21\"\n",
    "gene_id = \"ENSG00000269950\"\n",
    "z_cutoff_list = [2, 10, 20, 30, 40, 50]\n",
    "window_start = 0\n",
    "gene_window_size = 10000\n",
    "window_step_size = 200000\n",
    "window_max = 1000000\n",
    "maf_range_list = [[0, 0.01], [0.01, 0.05], [0.05, 0.1], [0.1, 0.5]]\n",
    "tpm_cutoff = 0.5\n",
    "\n",
    "ge_mtx_path = wkdir_path.joinpath(\"2_misexp_qc/misexp_gene_cov_corr/tpm_zscore_4568smpls_8610genes_flat_misexp_corr_qc.csv\")\n",
    "vcf_path = f\"/lustre/scratch126/humgen/projects/interval_wgs/final_release_freeze/gt_phased/interval_wgs.{chrom}.gt_phased.vcf.gz\"\n",
    "gt_info_root_path = wkdir_path.joinpath(\"4_vrnt_enrich/snp_indel_count_carriers/vrnts_gts_intersect\")\n",
    "wgs_rna_paired_smpls_path = wkdir_path.joinpath(\"1_rna_seq_qc/wgs_rna_match/paired_wgs_rna_postqc_prioritise_wgs.tsv\")\n",
    "genes_bed_path = wkdir_path.joinpath(f\"4_vrnt_enrich/snp_indel_count_carriers/genes_bed/{chrom}_genes.bed\")\n",
    "vrnts_bed_path = wkdir_path.joinpath(f\"4_vrnt_enrich/snp_indel_count_carriers/vrnts_bed/{chrom}_vrnts.bed\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "e3d7fbfe",
   "metadata": {},
   "outputs": [],
   "source": [
    "# carriers \n",
    "#vrnt_gt_intersect_path = Path(gt_info_root_path).joinpath(f\"{chrom}/{chrom}_{gene_id}_vrnts_gts_intersect.tsv\")\n",
    "#vrnt_gt_intersect_df = pd.read_csv(vrnt_gt_intersect_path, sep=\"\\t\")\n",
    "#vrnt_gt_intersect_df[\"gene_id\"] = gene_id"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "76b42d2e",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Counting variants upstream, 0bp ...\n",
      "Counting variants upstream, 200000bp ...\n",
      "Counting variants upstream, 400000bp ...\n",
      "Counting variants upstream, 600000bp ...\n",
      "Counting variants upstream, 800000bp ...\n",
      "Counting variants upstream, 1000000bp ...\n",
      "Counting variants downstream, 0bp ...\n",
      "Counting variants downstream, 200000bp ...\n",
      "Counting variants downstream, 400000bp ...\n",
      "Counting variants downstream, 600000bp ...\n",
      "Counting variants downstream, 800000bp ...\n",
      "Counting variants downstream, 1000000bp ...\n"
     ]
    }
   ],
   "source": [
    "### assign gene-variant pairs a window  \n",
    "vrnts_bed = BedTool(vrnts_bed_path)\n",
    "genes_bed = BedTool(genes_bed_path)\n",
    "\n",
    "\n",
    "intersect_bed_columns={0:\"chrom_gene\", 1:\"start_gene\", 2:\"end_gene\", 3: \"gene_id\", 4: \"score\", 5: \"strand\", \n",
    "                       6:\"chrom_vrnt\", 7:\"start_vrnt\", 8:\"end_vrnt\", 9:\"vrnt_id\"}\n",
    "vrnt_gene_pairs_in_windows_df_list = []\n",
    "for direction in [\"upstream\", \"downstream\"]:\n",
    "    window_size = window_start\n",
    "    seen_vrnt_gene_pairs = set()\n",
    "    while window_size <= window_max:\n",
    "        print(f\"Counting variants {direction}, {window_size}bp ...\")\n",
    "        window_name = f\"{direction}_{window_size}\"\n",
    "        if direction == \"upstream\": \n",
    "            l, r = window_size, 0 \n",
    "        else: \n",
    "            l, r = 0, window_size\n",
    "        vrnts_window_intersect_str = StringIO(str(genes_bed.window(vrnts_bed, \n",
    "                                                                   l=l, \n",
    "                                                                   r=r,\n",
    "                                                                   sw=True\n",
    "                                                                  )))        \n",
    "        intersect_bed_df = pd.read_csv(vrnts_window_intersect_str, sep=\"\\t\", header=None).rename(columns=intersect_bed_columns)\n",
    "        intersect_bed_df = intersect_bed_df[[\"vrnt_id\", \"gene_id\"]].copy()\n",
    "        # variant gene ID column \n",
    "        intersect_bed_df[\"vrnt_gene_id\"] = intersect_bed_df[\"vrnt_id\"].astype(str) + \",\" + intersect_bed_df[\"gene_id\"].astype(str)\n",
    "        # subset to variant gene pairs that have not been seen in previous windows    \n",
    "        vrnt_gene_pairs_in_window = set(intersect_bed_df.vrnt_gene_id.unique())\n",
    "        new_vrnt_gene_pairs = vrnt_gene_pairs_in_window - seen_vrnt_gene_pairs\n",
    "        new_vrnt_gene_pairs_df = intersect_bed_df[intersect_bed_df.vrnt_gene_id.isin(new_vrnt_gene_pairs)].copy()\n",
    "        new_vrnt_gene_pairs_df[\"window\"] = window_name\n",
    "        vrnt_gene_pairs_in_windows_df_list.append(new_vrnt_gene_pairs_df)\n",
    "        seen_vrnt_gene_pairs = seen_vrnt_gene_pairs.union(new_vrnt_gene_pairs)\n",
    "        window_size += window_step_size\n",
    "\n",
    "vrnt_gene_pairs_in_windows_df = pd.concat(vrnt_gene_pairs_in_windows_df_list)\n",
    "rename_window = {\"upstream_0\": \"gene_body\", \"downstream_0\": \"gene_body\"}\n",
    "vrnt_gene_pairs_in_windows_df.window = vrnt_gene_pairs_in_windows_df.window.replace(rename_window)\n",
    "vrnt_gene_pairs_in_windows_df = vrnt_gene_pairs_in_windows_df.drop_duplicates()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "9065a2ba",
   "metadata": {},
   "outputs": [],
   "source": [
    "# all possible SNP/indel carriers \n",
    "windows_maf_vrnt_type, count = {}, 0\n",
    "for window in vrnt_gene_pairs_in_windows_df.window.unique(): \n",
    "    for maf_bin in [\"0-1\", \"1-5\", \"5-10\", \"10-50\"]:\n",
    "        for vrnt_type in [\"snp\", \"indel\"]: \n",
    "            windows_maf_vrnt_type[count] = [window, maf_bin, vrnt_type]\n",
    "            count += 1\n",
    "windows_maf_vrnt_type_df = pd.DataFrame.from_dict(windows_maf_vrnt_type, orient=\"index\", columns=[\"window\", \"maf_bin\", \"vrnt_type\"])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "7d18f063",
   "metadata": {},
   "outputs": [],
   "source": [
    "## gene expression \n",
    "ge_mtx_df = pd.read_csv(ge_mtx_path)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "40c34bd2",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Number of misexpressed genes on chromosome: 82\n",
      "Samples with genotyping data and passing RNA-seq QC: 2821\n"
     ]
    }
   ],
   "source": [
    "### subset gene expresstion matrix to genes on chromsome \n",
    "misexp_genes = pd.read_csv(genes_bed_path, sep=\"\\t\", header=None)[3].unique()\n",
    "num_misexp_genes = len(misexp_genes)\n",
    "print(f\"Number of misexpressed genes on chromosome: {num_misexp_genes}\")\n",
    "ge_mtx_chrom_df = ge_mtx_df[ge_mtx_df.gene_id.isin(misexp_genes)]\n",
    "\n",
    "### subset to gene expression matrix to samples with genotyping data \n",
    "# get EGAN IDs with RNA data \n",
    "wgs_rna_paired_smpls_df = pd.read_csv(wgs_rna_paired_smpls_path, sep=\"\\t\")\n",
    "egan_ids_with_rna = wgs_rna_paired_smpls_df.egan_id.unique()\n",
    "# load vcf\n",
    "vcf = pysam.VariantFile(vcf_path, mode = \"r\")\n",
    "vcf_samples = [sample for sample in vcf.header.samples]\n",
    "vcf_egan_ids_with_rna = set(egan_ids_with_rna).intersection(set(vcf_samples))\n",
    "# subset egan ID and RNA ID links to samples with genotyping data and passing QC \n",
    "wgs_rna_paired_smpls_in_vcf_df = wgs_rna_paired_smpls_df[wgs_rna_paired_smpls_df.egan_id.isin(vcf_egan_ids_with_rna)]\n",
    "ge_matrix_flat_chrom_egan_df = pd.merge(ge_mtx_chrom_df, wgs_rna_paired_smpls_in_vcf_df, how=\"inner\", on=\"rna_id\")\n",
    "# add gene sample pairs \n",
    "ge_matrix_flat_chrom_egan_df[\"gene_smpl_pair\"] = ge_matrix_flat_chrom_egan_df.gene_id + \",\" + ge_matrix_flat_chrom_egan_df.egan_id\n",
    "\n",
    "rna_ids_pass_gt_rna_qc = ge_matrix_flat_chrom_egan_df.rna_id.unique()\n",
    "gene_ids_pass_qc = ge_matrix_flat_chrom_egan_df.gene_id.unique()\n",
    "print(f\"Samples with genotyping data and passing RNA-seq QC: {len(rna_ids_pass_gt_rna_qc)}\")\n",
    "if ge_matrix_flat_chrom_egan_df.shape[0] != num_misexp_genes * len(rna_ids_pass_gt_rna_qc): \n",
    "    raise ValueError(\"Number of genes and samples passsing QC does not match gene expression matrix size.\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "id": "80b0d34d",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Z-score threshold: 2\n",
      "\tNumber of misexpressed genes: 72\n",
      "\tENSG00000206102\n",
      "\tENSG00000224100\n",
      "\tENSG00000229289\n",
      "\tENSG00000229986\n",
      "\tENSG00000274248\n",
      "\tENSG00000237338\n",
      "\tENSG00000229382\n",
      "\tENSG00000142182\n",
      "\tENSG00000232360\n",
      "\tENSG00000182591\n",
      "\tENSG00000227702\n",
      "\tENSG00000236471\n",
      "\tENSG00000224541\n",
      "\tENSG00000233393\n",
      "\tENSG00000275874\n",
      "\tENSG00000184856\n",
      "\tENSG00000226956\n",
      "\tENSG00000225330\n",
      "\tENSG00000237735\n",
      "\tENSG00000205439\n",
      "\tENSG00000241123\n",
      "\tENSG00000280095\n",
      "\tENSG00000223400\n",
      "\tENSG00000274749\n",
      "\tENSG00000272804\n",
      "\tENSG00000188694\n",
      "\tENSG00000224574\n",
      "\tENSG00000230379\n",
      "\tENSG00000187766\n",
      "\tENSG00000231324\n",
      "\tENSG00000187026\n",
      "\tENSG00000267857\n",
      "\tENSG00000231620\n",
      "\tENSG00000227075\n",
      "\tENSG00000231986\n",
      "\tENSG00000231106\n",
      "\tENSG00000224269\n",
      "\tENSG00000206105\n",
      "\tENSG00000223608\n",
      "\tENSG00000237664\n",
      "\tENSG00000236545\n",
      "\tENSG00000259981\n",
      "\tENSG00000273115\n",
      "\tENSG00000198390\n",
      "\tENSG00000230978\n",
      "\tENSG00000261706\n",
      "\tENSG00000233480\n",
      "\tENSG00000277693\n",
      "\tENSG00000198054\n",
      "\tENSG00000236332\n",
      "\tENSG00000183640\n",
      "\tENSG00000224413\n",
      "\tENSG00000230794\n",
      "\tENSG00000231867\n",
      "\tENSG00000160202\n",
      "\tENSG00000235123\n",
      "\tENSG00000279579\n",
      "\tENSG00000223563\n",
      "\tENSG00000160181\n",
      "\tENSG00000160182\n",
      "\tENSG00000187175\n",
      "\tENSG00000227342\n",
      "\tENSG00000232806\n",
      "\tENSG00000186977\n",
      "\tENSG00000233316\n",
      "\tENSG00000184032\n",
      "\tENSG00000230972\n",
      "\tENSG00000221864\n",
      "\tENSG00000225637\n",
      "\tENSG00000281181\n",
      "\tENSG00000225298\n",
      "\tENSG00000226935\n"
     ]
    }
   ],
   "source": [
    "count_misexp_cntrl_carriers_windows_df_list = []\n",
    "for z_cutoff in z_cutoff_list:\n",
    "    # misexpression events \n",
    "    print(f\"Z-score threshold: {z_cutoff}\")\n",
    "    misexp_df = ge_matrix_flat_chrom_egan_df[(ge_matrix_flat_chrom_egan_df[\"z-score\"] > z_cutoff) & \n",
    "                                             (ge_matrix_flat_chrom_egan_df[\"TPM\"] > tpm_cutoff)]\n",
    "    misexp_genes = misexp_df.gene_id.unique()\n",
    "    print(f\"\\tNumber of misexpressed genes: {len(misexp_genes)}\")\n",
    "    for gene_id in misexp_genes:\n",
    "        print(f\"\\t{gene_id}\")\n",
    "        gene_misexp_df = misexp_df[misexp_df.gene_id == gene_id]\n",
    "        # misexpressed samples \n",
    "        misexp_rna_id = gene_misexp_df.rna_id.unique()\n",
    "        misexp_gene_smpl = gene_misexp_df.gene_smpl_pair.unique()\n",
    "        # control events\n",
    "        cntrl_df = ge_matrix_flat_chrom_egan_df[(ge_matrix_flat_chrom_egan_df.gene_id == gene_id) & \n",
    "                                                (~ge_matrix_flat_chrom_egan_df.gene_smpl_pair.isin(misexp_gene_smpl))]\n",
    "        cntrl_gene_smpl = cntrl_df.gene_smpl_pair.unique() \n",
    "        # load genotypes \n",
    "        vrnt_gt_intersect_path = Path(gt_info_root_path).joinpath(f\"{chrom}/{chrom}_{gene_id}_vrnts_gts_intersect.tsv\")\n",
    "        vrnt_gt_intersect_df = pd.read_csv(vrnt_gt_intersect_path, sep=\"\\t\")\n",
    "        vrnt_gt_intersect_df[\"gene_id\"] = gene_id\n",
    "        # assign variants MAF bins \n",
    "        maf_bins = [0, 0.01, 0.05, 0.1, 0.5]\n",
    "        af_bin_labels = [\"0-1\", \"1-5\", \"5-10\", \"10-50\"]\n",
    "        vrnt_gt_maf_lt50_df = vrnt_gt_intersect_df[vrnt_gt_intersect_df.AF < 0.5].copy()\n",
    "        vrnt_gt_maf_lt50_df[\"maf_bin\"] = pd.cut(vrnt_gt_maf_lt50_df.AF, bins=maf_bins, labels=af_bin_labels, right=False)\n",
    "        maf_bins = [1-0.5, 1-0.1, 1-0.05, 1-0.01, 1-0]\n",
    "        af_bin_labels =[\"10-50\", \"5-10\", \"1-5\", \"0-1\"]\n",
    "        vrnt_gt_maf_gt50_df = vrnt_gt_intersect_df[vrnt_gt_intersect_df.AF >= 0.5].copy()\n",
    "        vrnt_gt_maf_gt50_df[\"maf_bin\"] = pd.cut(vrnt_gt_maf_gt50_df.AF, bins=maf_bins, labels=af_bin_labels, right=True)\n",
    "        vrnt_gt_maf_df = pd.concat([vrnt_gt_maf_gt50_df, vrnt_gt_maf_lt50_df])\n",
    "        # add position windows \n",
    "        vrnt_gt_af_pos_df = pd.merge(vrnt_gt_maf_df, \n",
    "                                     vrnt_gene_pairs_in_windows_df, \n",
    "                                     on=[\"vrnt_id\", \"gene_id\"], \n",
    "                                     how=\"inner\")\n",
    "        # add gene-sample IDs \n",
    "        vrnt_gt_af_pos_df[\"gene_smpl_pair\"] = vrnt_gt_af_pos_df.gene_id + \",\" + vrnt_gt_af_pos_df.egan_id\n",
    "        # misexp gene-pairs \n",
    "        misexp_carriers_df = vrnt_gt_af_pos_df[vrnt_gt_af_pos_df.gene_smpl_pair.isin(misexp_gene_smpl)]\n",
    "        \n",
    "        count_misexp_carriers_df = pd.DataFrame(misexp_carriers_df.groupby([\"window\", \"vrnt_type\", \"maf_bin\"]).gene_smpl_pair.nunique())\n",
    "        count_misexp_carriers_df = count_misexp_carriers_df.reset_index()\n",
    "        count_misexp_carriers_df = count_misexp_carriers_df.rename(columns={\"gene_smpl_pair\": \"misexp_carrier_check\"})\n",
    "        \n",
    "        # cntrl gene_pairs \n",
    "        cntrl_carriers_df = vrnt_gt_af_pos_df[vrnt_gt_af_pos_df.gene_smpl_pair.isin(cntrl_gene_smpl)]\n",
    "        count_cntrl_carriers_df = pd.DataFrame(cntrl_carriers_df.groupby([\"window\", \"vrnt_type\", \"maf_bin\"]).gene_smpl_pair.nunique())\n",
    "        count_cntrl_carriers_df = count_cntrl_carriers_df.reset_index()\n",
    "        count_cntrl_carriers_df = count_cntrl_carriers_df.rename(columns={\"gene_smpl_pair\": \"control_carrier_check\"})\n",
    "        \n",
    "        count_misexp_cntrl_carrier_df = pd.merge(count_misexp_carriers_df, \n",
    "                                                 count_cntrl_carriers_df, \n",
    "                                                 on=[\"window\", \"maf_bin\", \"vrnt_type\"], \n",
    "                                                 how=\"outer\"\n",
    "                                                )\n",
    "        count_misexp_cntrl_carrier_all_df = pd.merge(windows_maf_vrnt_type_df, \n",
    "                                                     count_misexp_cntrl_carrier_df, \n",
    "                                                     on=[\"window\", \"maf_bin\", \"vrnt_type\"], \n",
    "                                                     how=\"left\")\n",
    "        count_misexp_cntrl_carrier_all_df[\"gene_id\"] = gene_id\n",
    "        count_misexp_cntrl_carrier_all_df[\"misexp_total_check\"] = len(misexp_gene_smpl)\n",
    "        count_misexp_cntrl_carrier_all_df[\"control_total_check\"] = len(cntrl_gene_smpl)\n",
    "        count_misexp_cntrl_carrier_all_df[\"z_cutoff\"] = z_cutoff\n",
    "        count_misexp_cntrl_carriers_windows_df_list.append(count_misexp_cntrl_carrier_all_df)\n",
    "    break \n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "id": "8cca077c",
   "metadata": {},
   "outputs": [],
   "source": [
    "count_misexp_cntrl_carriers_windows_df = pd.concat(count_misexp_cntrl_carriers_windows_df_list) \n",
    "count_misexp_cntrl_carriers_windows_df.maf_bin = count_misexp_cntrl_carriers_windows_df.maf_bin.astype(str)\n",
    "count_misexp_cntrl_carriers_windows_df = count_misexp_cntrl_carriers_windows_df.fillna(0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "id": "12056018",
   "metadata": {},
   "outputs": [],
   "source": [
    "### load previous results \n",
    "snp_indel_carriers_results_path = wkdir_path.joinpath(f\"4_vrnt_enrich/snp_indel_count_carriers/count_snp_indel_carriers_af50/{chrom}_carrier_count.tsv\")\n",
    "snp_indel_carriers_results_df = pd.read_csv(snp_indel_carriers_results_path, sep=\"\\t\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "id": "75b5b70a",
   "metadata": {},
   "outputs": [],
   "source": [
    "rename_maf_range = {'0-0.01': \"0-1\", '0.01-0.05': \"1-5\", '0.05-0.1': \"5-10\", '0.1-0.5': \"10-50\"}\n",
    "snp_indel_carriers_results_df[\"maf_bin\"] = snp_indel_carriers_results_df.maf_range.replace(rename_maf_range)\n",
    "snp_indel_carriers_results_df = snp_indel_carriers_results_df.drop(columns=[\"maf_range\"])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "id": "94e6d5f6",
   "metadata": {},
   "outputs": [],
   "source": [
    "snp_indel_count_window_check_df = pd.merge(snp_indel_carriers_results_df, \n",
    "                                    count_misexp_cntrl_carriers_windows_df, \n",
    "                                    on=[\"z_cutoff\", \"window\", \"gene_id\", \"vrnt_type\", \"maf_bin\"],\n",
    "                                    how = \"inner\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "id": "5136381d",
   "metadata": {},
   "outputs": [],
   "source": [
    "# check carrier counts\n",
    "if not snp_indel_count_window_check_df.control_total.astype(int).equals(snp_indel_count_window_check_df.control_total_check.astype(int)):\n",
    "    raise ValueError(\"Difference in total control events.\")\n",
    "if not snp_indel_count_window_check_df.misexp_total.astype(int).equals(snp_indel_count_window_check_df.misexp_total_check.astype(int)):\n",
    "    raise ValueError(\"Difference in total misexpression events.\")\n",
    "if not snp_indel_count_window_check_df.misexp_carrier.astype(int).equals(snp_indel_count_window_check_df.misexp_carrier_check.astype(int)):\n",
    "    raise ValueError(\"Difference in total misexpression carriers.\")\n",
    "if not snp_indel_count_window_check_df.control_carrier.astype(int).equals(snp_indel_count_window_check_df.control_carrier_check.astype(int)):\n",
    "    raise ValueError(\"Difference in total control carriers.\")"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
